package com.alibaba.ttl.threadpool;

import cn.micro.ttl.mdc.MDCThreadWrap;
import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import com.alibaba.ttl.spi.TtlEnhanced;
import com.alibaba.ttl.spi.TtlWrapper;
import org.springframework.scheduling.SchedulingTaskExecutor;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.Trigger;

import java.util.Date;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledFuture;

public class TaskSchedulerTtlMDCWrapper implements TtlWrapper<TaskScheduler>, TtlEnhanced, TaskScheduler, SchedulingTaskExecutor {

    private final TaskScheduler taskScheduler;
    private final SchedulingTaskExecutor schedulingTaskExecutor;

    protected final boolean idempotent;

    public TaskSchedulerTtlMDCWrapper(TaskScheduler taskScheduler, SchedulingTaskExecutor schedulingTaskExecutor, boolean idempotent) {
        this.taskScheduler = taskScheduler;
        this.schedulingTaskExecutor = schedulingTaskExecutor;
        this.idempotent = idempotent;
    }

    @Override
    public TaskScheduler unwrap() {
        return taskScheduler;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;

        TaskSchedulerTtlMDCWrapper that = (TaskSchedulerTtlMDCWrapper) o;

        return taskScheduler.equals(that.taskScheduler);
    }

    @Override
    public int hashCode() {
        return taskScheduler.hashCode();
    }

    @Override
    public String toString() {
        return this.getClass().getName() + " - " + taskScheduler.toString();
    }

    @Override
    public ScheduledFuture<?> schedule(Runnable runnable, Trigger trigger) {
        return taskScheduler.schedule(TtlRunnable.get(MDCThreadWrap.runnable(runnable), false, idempotent), trigger);
    }

    @Override
    public ScheduledFuture<?> schedule(Runnable runnable, Date date) {
        return taskScheduler.schedule(TtlRunnable.get(MDCThreadWrap.runnable(runnable), false, idempotent), date);
    }

    @Override
    public ScheduledFuture<?> scheduleAtFixedRate(Runnable runnable, Date date, long l) {
        return taskScheduler.scheduleAtFixedRate(TtlRunnable.get(MDCThreadWrap.runnable(runnable), false, idempotent), date, l);

    }

    @Override
    public ScheduledFuture<?> scheduleAtFixedRate(Runnable runnable, long l) {
        return taskScheduler.scheduleAtFixedRate(TtlRunnable.get(MDCThreadWrap.runnable(runnable), false, idempotent), l);

    }

    @Override
    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable runnable, Date date, long l) {
        return taskScheduler.scheduleWithFixedDelay(TtlRunnable.get(MDCThreadWrap.runnable(runnable), false, idempotent), date, l);

    }

    @Override
    public ScheduledFuture<?> scheduleWithFixedDelay(Runnable runnable, long l) {
        return taskScheduler.scheduleWithFixedDelay(TtlRunnable.get(MDCThreadWrap.runnable(runnable), false, idempotent), l);

    }

    @Override
    public void execute(Runnable task, long startTimeout) {
        schedulingTaskExecutor.execute(TtlRunnable.get(MDCThreadWrap.runnable(task), false, idempotent), startTimeout);
    }

    @Override
    public Future<?> submit(Runnable task) {
        return schedulingTaskExecutor.submit(TtlRunnable.get(MDCThreadWrap.runnable(task), false, idempotent));

    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        return schedulingTaskExecutor.submit(TtlCallable.get(MDCThreadWrap.callable(task), false, idempotent));

    }

    @Override
    public void execute(Runnable task) {
        schedulingTaskExecutor.execute(TtlRunnable.get(MDCThreadWrap.runnable(task), false, idempotent));
    }
}
